import torch as t
from scipy import ndimage
from skimage.feature import canny
import numpy as np
import scipy
import cv2
from scipy.ndimage import gaussian_filter

class conv_detector(object):
    def __init__(self):
        pass
    
    def run(self,kernel,img,thre=0.2,sigma=0.1,progresser=None,outputer=None):
        self.kernel = kernel
        self.img = img
        self.outputer = outputer
        kernel_tensor = t.tensor(kernel)
        img_tensor = t.tensor(img)
        if type(progresser) != type(None):
            progresser.value = 10
        conved_img = self.conv_img(img_tensor,kernel_tensor).detach().cpu().numpy()
        #conved_img = np.clip(conved_img,np.max(conved_img)*thre,np.max(conved_img))
        
        if type(progresser) != type(None):
            progresser.value = 60
        conved_img = gaussian_filter(conved_img, sigma=sigma*100)
        cor = self.find_local_maxima(conved_img)

        if type(progresser) != type(None):
            progresser.value = 80
        list_0,list_1,list_2 = self.classify_coor(cor)
        tree = scipy.spatial.cKDTree(cor)
        neigh = tree.query_ball_point(list_1[0], r=40)
        with outputer:
            print('Gap = ',cor[neigh[0],:]-list_1[0],' pixels')

        img_rgb = self.get_rgb(list_0,list_1,list_2)
        return img_rgb
    

    def get_rgb(self,list_0,list_1,list_2):
        img_rgb = np.zeros((self.img.shape[0],self.img.shape[1],3))
        img_rgb[:,:,2] = self.img/np.max(self.img)
        img_g = np.zeros((self.img.shape[0]+self.kernel.shape[0],self.img.shape[1]+self.kernel.shape[1]))
        img_b = np.zeros((self.img.shape[0]+self.kernel.shape[0],self.img.shape[1]+self.kernel.shape[1]))
        for item in list_1:
            x1 = int(item[0])
            x2 = int(item[0])+self.kernel.shape[0]
            y1 = int(item[1])
            y2 = int(item[1])+self.kernel.shape[1]
            img_g[x1:x2,y1:y2] += self.kernel
            
        for item in list_2:
            x1 = int(item[0])
            x2 = int(item[0])+self.kernel.shape[0]
            y1 = int(item[1])
            y2 = int(item[1])+self.kernel.shape[1]
            img_b[x1:x2,y1:y2] += self.kernel
        img_rgb[:,:,0] = img_g[:self.img.shape[0],
                                :self.img.shape[1]]
        img_rgb[:,:,1] = img_b[:self.img.shape[0],
                                :self.img.shape[1]]
        # img_rgb[:,:,0] = img_g[(self.kernel.shape[0]-1)//2:(self.kernel.shape[0]-1)//2+self.img.shape[0],
        #                         (self.kernel.shape[1]-1)//2:(self.kernel.shape[1]-1)//2+self.img.shape[1]]
        # img_rgb[:,:,1] = img_b[(self.kernel.shape[0]-1)//2:(self.kernel.shape[0]-1)//2+self.img.shape[0],
        #                         (self.kernel.shape[1]-1)//2:(self.kernel.shape[1]-1)//2+self.img.shape[1]]
        img_rgb[:,:,0] /= np.max(img_rgb[:,:,0])
        img_rgb[:,:,1] /= np.max(img_rgb[:,:,1])
        img_rgb[:,:,:2][img_rgb[:,:,:2]>0.1] = 1
        return img_rgb

    def classify_coor(self,cor):
        tree = scipy.spatial.cKDTree(cor)
        list_0 = []
        list_1 = []
        list_2 = []
        for center in range(cor.shape[0]):
            cor_center = cor[center,:]
            neigh = tree.query_ball_point(cor_center, r=40)
            if len(neigh) > 1:
                if len(list_1) == 0:
                    list_1.append(cor_center)
                else:
                    neigh = tree.query(cor_center,k=4)
                    neigh1 = tree.query(list_1[0],k=4)
                    if np.mean((cor[neigh[1],:][1,:]-cor[neigh1[1],:][1,:]-cor_center+list_1[0])**2)<1000:
                        list_1.append(cor_center)
                    else:
                        list_2.append(cor_center)
            else:
                list_0.append(cor_center)
        return list_0,list_1,list_2


    def get_coor(self,img,sigma=5):
        img = img/np.max(img)
        edges = canny(img, sigma=sigma)
        fills = ndimage.binary_fill_holes(edges)
        labels, nlabels = ndimage.label(fills,structure=[[0,1,0],
        [1,1,1],
        [0,1,0]])
        coordinates = np.array(
        ndimage.center_of_mass(
            fills, labels, np.arange(nlabels) + 1))
        return coordinates


    def conv_img(self,img,kernel,stride=(1,1)):
        img_extend = img.unsqueeze(0).unsqueeze(0).float()
        kernel_extend = kernel.unsqueeze(0).unsqueeze(0).float()
        c = t.nn.Conv2d(1,1,kernel_extend.shape,stride=stride,padding=((kernel_extend.shape[0]-1)//2,(kernel_extend.shape[1]-1)//2),bias=False)
        c.weight.data = kernel_extend
        return c(img_extend).squeeze(0).squeeze(0)


    def find_local_maxima(img):
        # 定义结构元素
        kernel = np.ones((40, 40), np.uint8)
        # img = gaussian_filter(img, sigma=2)
        # 膨胀操作
        dilated = cv2.dilate(img, kernel, iterations=1)
        # 提取局部极大值
        max_indices = np.where(dilated == img)
        # 返回结果
        return np.array(max_indices).T




